import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import numpy as np
import random
from torch.utils.data import Dataset, DataLoader
import math

# 设置随机种子
seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(seed)

# 设备配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 1. 位置编码
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

# 2. 编码器层
class EncoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super(EncoderLayer, self).__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        
        self.activation = F.relu
        
    def forward(self, src, src_mask=None, src_key_padding_mask=None):
        src2 = self.self_attn(src, src, src, attn_mask=src_mask,
                              key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src

# 3. 解码器层
class DecoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super(DecoderLayer, self).__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)
        
        self.activation = F.relu
        
    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None,
                tgt_key_padding_mask=None, memory_key_padding_mask=None):
        tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)
        
        tgt2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)
        
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)
        return tgt

# 4. 完整Transformer模型
class TransformerModel(nn.Module):
    def __init__(self, num_tokens, d_model=512, nhead=8, num_encoder_layers=6,
                 num_decoder_layers=6, dim_feedforward=2048, dropout=0.1):
        super(TransformerModel, self).__init__()
        
        self.embedding = nn.Embedding(num_tokens, d_model)
        self.d_model = d_model
        self.pos_encoder = PositionalEncoding(d_model, dropout=dropout)
        
        encoder_layers = EncoderLayer(d_model, nhead, dim_feedforward, dropout)
        self.encoder = nn.TransformerEncoder(encoder_layers, num_layers=num_encoder_layers)
        
        decoder_layers = DecoderLayer(d_model, nhead, dim_feedforward, dropout)
        self.decoder = nn.TransformerDecoder(decoder_layers, num_layers=num_decoder_layers)
        
        self.fc_out = nn.Linear(d_model, num_tokens)
        self.init_weights()
        
    def init_weights(self):
        initrange = 0.1
        self.embedding.weight.data.uniform_(-initrange, initrange)
        self.fc_out.bias.data.zero_()
        self.fc_out.weight.data.uniform_(-initrange, initrange)
        
    def forward(self, src, tgt, src_mask=None, tgt_mask=None,
                memory_mask=None, src_key_padding_mask=None,
                tgt_key_padding_mask=None, memory_key_padding_mask=None):
        src = self.embedding(src) * math.sqrt(self.d_model)
        src = self.pos_encoder(src)
        
        tgt = self.embedding(tgt) * math.sqrt(self.d_model)
        tgt = self.pos_encoder(tgt)
        
        memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask)
        output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
                              tgt_key_padding_mask=tgt_key_padding_mask,
                              memory_key_padding_mask=memory_key_padding_mask)
        
        output = self.fc_out(output)
        return output

# 5. 数据集
class SequenceDataset(Dataset):
    def __init__(self, num_samples=10000, min_len=3, max_len=6, max_value=9):
        self.num_samples = num_samples
        self.min_len = min_len
        self.max_len = max_len
        self.max_value = max_value
        self.data = self.generate_data()
        
    def generate_data(self):
        data = []
        for _ in range(self.num_samples):
            seq_len = random.randint(self.min_len, self.max_len)
            seq = [random.randint(1, self.max_value) for _ in range(seq_len)]
            data.append(seq)
        return data
    
    def __len__(self):
        return self.num_samples
    
    def __getitem__(self, idx):
        sequence = self.data[idx]
        return torch.tensor(sequence, dtype=torch.long)

# 数据预处理
def preprocess_sequence(sequence, start_token=0, end_token=10):
    encoder_input = sequence.clone()
    decoder_input = torch.cat([torch.tensor([start_token]), sequence])
    target = torch.cat([sequence, torch.tensor([end_token])])
    return encoder_input, decoder_input, target

# 批次处理（填充）
def collate_fn(batch, start_token=0, end_token=10, pad_token=11):
    processed = [preprocess_sequence(seq, start_token, end_token) for seq in batch]
    encoder_inputs, decoder_inputs, targets = zip(*processed)
    
    max_enc_len = max(len(seq) for seq in encoder_inputs)
    max_dec_len = max(len(seq) for seq in decoder_inputs)
    max_target_len = max(len(seq) for seq in targets)
    
    padded_enc = []
    padded_dec = []
    padded_targets = []
    
    for enc, dec, target in zip(encoder_inputs, decoder_inputs, targets):
        pad_enc = torch.cat([enc, torch.full((max_enc_len - len(enc),), pad_token, dtype=torch.long)])
        padded_enc.append(pad_enc)
        
        pad_dec = torch.cat([dec, torch.full((max_dec_len - len(dec),), pad_token, dtype=torch.long)])
        padded_dec.append(pad_dec)
        
        pad_target = torch.cat([target, torch.full((max_target_len - len(target),), pad_token, dtype=torch.long)])
        padded_targets.append(pad_target)
    
    return (
        torch.stack(padded_enc),  # 形状: (batch_size, max_enc_len)
        torch.stack(padded_dec),  # 形状: (batch_size, max_dec_len)
        torch.stack(padded_targets)  # 形状: (batch_size, max_target_len)
    )

# 6. 掩码生成函数（核心修正）
def generate_mask(src, tgt, pad_token=11):
    """
    生成符合PyTorch要求的掩码：
    - src_pad_mask: (batch_size, src_seq_len)
    - tgt_pad_mask: (batch_size, tgt_seq_len)
    - tgt_mask: (tgt_seq_len, tgt_seq_len)
    """
    src_seq_len = src.size(1)  # src形状: (batch_size, src_seq_len)
    tgt_seq_len = tgt.size(1)  # tgt形状: (batch_size, tgt_seq_len)
    
    # 解码器自注意力掩码：(tgt_seq_len, tgt_seq_len)
    tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt_seq_len).to(device)
    
    # 编码器自注意力掩码：通常为None
    src_mask = None
    
    # 填充掩码：(batch_size, seq_len)，无需转置！
    src_pad_mask = (src == pad_token)  # 直接比较，形状正确
    tgt_pad_mask = (tgt == pad_token)  # 直接比较，形状正确
    
    return src_mask, tgt_mask, src_pad_mask, tgt_pad_mask

# 7. 训练函数
def train_model(model, train_loader, criterion, optimizer, num_epochs=10, 
                start_token=0, end_token=10, pad_token=11):
    model.train()
    for epoch in range(num_epochs):
        total_loss = 0
        for batch in train_loader:
            encoder_inputs, decoder_inputs, targets = batch  # 形状: (batch_size, seq_len)
            
            # 转置为Transformer要求的(seq_len, batch_size)格式（仅输入需要，掩码不需要）
            encoder_inputs = encoder_inputs.transpose(0, 1).to(device)  # (src_seq_len, batch_size)
            decoder_inputs = decoder_inputs.transpose(0, 1).to(device)  # (tgt_seq_len, batch_size)
            targets = targets.transpose(0, 1).to(device)  # (tgt_seq_len, batch_size)
            
            # 生成掩码（传入原始形状的src和tgt：(batch_size, seq_len)）
            src_mask, tgt_mask, src_pad_mask, tgt_pad_mask = generate_mask(
                batch[0],  # encoder_inputs原始形状: (batch_size, src_seq_len)
                batch[1],  # decoder_inputs原始形状: (batch_size, tgt_seq_len)
                pad_token=pad_token
            )
            # 将掩码移到设备
            src_pad_mask = src_pad_mask.to(device)
            tgt_pad_mask = tgt_pad_mask.to(device)
            
            # 前向传播
            optimizer.zero_grad()
            outputs = model(
                encoder_inputs, 
                decoder_inputs,
                src_mask=src_mask,
                tgt_mask=tgt_mask,
                src_key_padding_mask=src_pad_mask,  # 形状: (batch_size, src_seq_len)
                tgt_key_padding_mask=tgt_pad_mask,  # 形状: (batch_size, tgt_seq_len)
                memory_key_padding_mask=src_pad_mask  # 编码器输出的填充掩码
            )
            
            # 计算损失
            loss = criterion(
                outputs.reshape(-1, outputs.size(-1)), 
                targets.reshape(-1)
            )
            
            # 反向传播
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
        
        avg_loss = total_loss / len(train_loader)
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {avg_loss:.4f}')
        
        if (epoch+1) % 5 == 0:
            test_sequence = [1, 2, 3, 4, 5]
            predicted = predict_sequence(model, test_sequence, start_token, end_token, pad_token)
            print(f'Test sequence: {test_sequence}')
            print(f'Predicted sequence: {predicted}')

# 8. 预测函数
def predict_sequence(model, input_sequence, start_token=0, end_token=10, pad_token=11, max_len=20):
    model.eval()
    with torch.no_grad():
        # 编码器输入: (1, seq_len) -> 转置为(seq_len, 1)
        encoder_input = torch.tensor(input_sequence, dtype=torch.long).unsqueeze(0).to(device)  # (1, seq_len)
        encoder_input = encoder_input.transpose(0, 1)  # (seq_len, 1)
        
        # 解码器初始输入: (1, 1) -> 转置为(1, 1)
        decoder_input = torch.tensor([[start_token]], dtype=torch.long).to(device)  # (1, 1)
        decoder_input = decoder_input.transpose(0, 1)  # (1, 1)
        
        # 编码器前向传播
        src = model.embedding(encoder_input) * math.sqrt(model.d_model)
        src = model.pos_encoder(src)
        memory = model.encoder(src)
        
        predicted_sequence = []
        
        for _ in range(max_len):
            tgt_mask = nn.Transformer.generate_square_subsequent_mask(decoder_input.size(0)).to(device)
            
            tgt = model.embedding(decoder_input) * math.sqrt(model.d_model)
            tgt = model.pos_encoder(tgt)
            output = model.decoder(tgt, memory, tgt_mask=tgt_mask)
            output = model.fc_out(output)
            
            predicted_token = output[-1, :, :].argmax(dim=1).item()
            
            if predicted_token == end_token or predicted_token == pad_token:
                break
            
            predicted_sequence.append(predicted_token)
            
            new_token = torch.tensor([[predicted_token]], dtype=torch.long).to(device)
            decoder_input = torch.cat([decoder_input, new_token.transpose(0, 1)], dim=0)
        
        return predicted_sequence

# 9. 主函数
def main():
    pad_token = 11
    start_token = 0
    end_token = 10
    num_tokens = 12  # 0-11
    d_model = 64
    nhead = 2
    num_encoder_layers = 2
    num_decoder_layers = 2
    dim_feedforward = 128
    dropout = 0.1
    batch_size = 32
    num_epochs = 20
    learning_rate = 0.001
    
    dataset = SequenceDataset(num_samples=10000, min_len=3, max_len=6)
    train_loader = DataLoader(
        dataset, 
        batch_size=batch_size, 
        shuffle=True, 
        collate_fn=lambda x: collate_fn(x, start_token, end_token, pad_token)
    )
    
    model = TransformerModel(
        num_tokens=num_tokens,
        d_model=d_model,
        nhead=nhead,
        num_encoder_layers=num_encoder_layers,
        num_decoder_layers=num_decoder_layers,
        dim_feedforward=dim_feedforward,
        dropout=dropout
    ).to(device)

    # model.encoder()
    
    criterion = nn.CrossEntropyLoss(ignore_index=pad_token)
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    
    print(f"Using device: {device}")
    train_model(model, train_loader, criterion, optimizer, num_epochs, 
                start_token, end_token, pad_token)
    
    # 保存模型
    torch.save(model.encoder, './encoder.pth')
    torch.save(model.decoder, './decoder.pth')
    torch.save(model.embedding, './embedding.pth')
    torch.save(model.fc_out, './fc_out.pth')
    print("模型保存完成")
    
    # 测试
    test_sequences = [
        [1, 2, 3, 4, 5],
        [3, 1, 4, 1, 5],
        [9, 8, 7]
    ]
    
    for seq in test_sequences:
        predicted = predict_sequence(model, seq, start_token, end_token, pad_token)
        print(f"输入序列: {seq}")
        print(f"预测序列: {predicted}\n")

if __name__ == "__main__":
    main()
    